'''
1.0.0 版本 
# Date: 2024/12/28
# Author: ZHENG
# Description:
    # 此版本功能如下：
        1. 通讯客户端,与ROS服务器建立连接,接收数据并解析,处理数据并更新UI

'''

# -*- coding: utf-8 -*-
import socket
import json
import io
import os
import struct
from threading import Thread
import threading
import tkinter as tk
from tkinter import messagebox


class ROS_Client:

    def __init__(self, server_ip = '', server_port = None):
        self.server_ip = server_ip
        self.server_port = server_port
        self.client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.client_connected = False
        self.pmg_flag = False

        self.buffer = b''

        self.ros_battery = None

        self.ros_speed = 0
        self.ros_angle = 0
        self.ros_mirobotstate = None
        self.ros_astra = None
        self.ros_emergency = None
        self.ros_controlMode = None

        self.ros_imuData = None
        self.ros_odomData = None
        self.ros_batteryData = None

    def connect(self):
        try:
            # self.client_socket.settimeout(5)
            self.client_socket.connect((self.server_ip, self.server_port))
            self.client_connected = True
            self.show_popup("TCP连接状态", "成功连接上ROS服务器")

            self.recv_thread = Thread(target=self.receive_data)
            self.recv_thread.start()

        # except (socket.timeout, socket.error) as e:
        except Exception as e:
            print(f"Failed to connect: {e}")
            self.show_popup("TCP连接状态", f"连接失败，原因: {e}")
            return False

    def receive_data(self):
        while True:
            try:
                recv_data = self.client_socket.recv(4096)
                # data = recv_data.decode('utf-8')
                data = recv_data
                if data is  None:
                    break
                if self.pmg_flag:
                    # self.parse_pgmData(data)
                    return data
                    # print(data.decode('utf-8'))
                else:
                    # print(data)
                    self.parse_data(data)
                # print(data.decode('utf-8'))
            
            except Exception as e:
                print(f"receive_data Error: {e}")
                continue

    ### 方向控制----start
    def send_direct_command(self, command):
        if self.client_connected:
            try:
                self.client_socket.sendall(command)
                return True
            
            except Exception as e:
                print(f"Send command failed: {e}")
                return False
        return False

    def start_continuous(self, linear_x=0, linear_y=0, angular_z=0):
        """启动持续运动"""
        self.speed_commands = {
            "LINEAR_X": linear_x,
            "LINEAR_Y": linear_y,
            "ANGULAR_Z": angular_z
        }
        if self.timer is None:
            self.timer = threading.Timer(0.1, self._send_continuous)
            self.timer.daemon = True
            self.timer.start()

    def stop_continuous(self):
        """停止持续运动"""
        if self.timer:
            self.timer.cancel()
            self.timer = None
        self.send_command("STOP")

    def _send_continuous(self):
        if self.client_connected:
            cmd = f"CONTINUOUS {self.speed_commands['LINEAR_X']} {self.speed_commands['LINEAR_Y']} {self.speed_commands['ANGULAR_Z']}"
            self.send_command(cmd)
            self.timer = threading.Timer(0.1, self._send_continuous)
            self.timer.daemon = True
            self.timer.start()

    ### 方向控制----end

    def send_command(self, controller= "ZBot3", ros_core = None, ros_run = None, ros_launch = None, ros_topic = None, ros_service = None, ros_gcode = None, ros_bag = None, data = None):
        try:
            message_body = self.create_message_body(controller, ros_core, ros_run, ros_launch, ros_topic, ros_service, ros_gcode, ros_bag, data)
            packet = self.build_packet(message_body)
            self.client_socket.sendall(packet)
        
        except Exception as e:
            print(f"ROS_Client send_command Error: {e}")

    def create_message_body(self, controller, ros_core, ros_run, ros_launch, ros_topic, ros_service, ros_gcode, ros_bag, data):
        message_body = { 
            "ZBot":controller,
            "cmdType":{
                "roscore": ros_core,
                "rosrun": ros_run,
                "roslaunch":ros_launch,
                "topic": ros_topic,
                "service": ros_service,
                "GCode": ros_gcode,
            },
            "rosBag":ros_bag,
            "cmdData": data,
        }
        return json.dumps(message_body) + '\n'

    def build_packet(self, message_body):
        header = json.dumps({  
            "data_size": len(message_body),  
            "type": "ros"  
        }).encode('utf-8')

        frame_header = len(header).to_bytes(4, 'little')
        packet = frame_header + header + message_body.encode('utf-8')
        return packet

    def close(self):
        self.client_socket.close()

    ### 调用显示提示对话框窗口
    def show_popup(self, title, message):
        root = tk.Tk()
        root.withdraw()
        messagebox.showinfo(title, message)
        root.destroy()

    def parse_pgmData(self, recv_data, save_path='maps/default.pgm'):
        try:
            os.makedirs(os.path.dirname(save_path), exist_ok=True)
            image_file = io.BytesIO(recv_data)
            with open(save_path, 'wb') as f:
                f.write(image_file.read())

        except Exception as e:
            print(f"ROS_Client parse_pgmData Error: {e}")


    def parse_data(self, new_data):
        self.buffer += new_data
        processed_data = b''
        
        while len(self.buffer) >= 4:
            try:
                # 1. 解析头部长度
                header_length = struct.unpack('I', self.buffer[:4])[0]
                
                # 检查是否收到完整头部
                if len(self.buffer) < 4 + header_length:
                    break
                
                # 2. 解析头部JSON
                header_data = self.buffer[4:4+header_length].decode()
                header = json.loads(header_data)
                data_size = header["data_size"]
                
                # 检查是否收到完整消息体
                total_length = 4 + header_length + data_size
                if len(self.buffer) < total_length:
                    break
                
                # 3. 解析消息体JSON
                message_body = self.buffer[4+header_length:total_length].decode()
                recv_str = json.loads(message_body)
                
                # 4. 处理传感器数据
                # print(f"recv_str {recv_str}")
                self.process_sensor_data(recv_str)
                
                # 5. 从缓冲区移除已处理数据
                self.buffer = self.buffer[total_length:]
                processed_data = self.buffer
                
            except (struct.error, KeyError, json.JSONDecodeError) as e:
                print(f"数据解析错误: {e}")
                # 跳过错误数据，寻找下一个有效包头
                if len(self.buffer) > 4:
                    self.buffer = self.buffer[1:]  # 滑动窗口寻找下一个包头
                else:
                    self.buffer = b''
                    break
        
        return processed_data

    def process_sensor_data(self, recv_str):
        """处理解析后的传感器数据"""
        if 'sensorData_' in recv_str: 
            sensor_data = recv_str["sensorData_"]
            self.ros_battery = sensor_data.get('battery')
            self.ros_speed = sensor_data.get('speed')
            self.ros_angle = sensor_data.get('angle')
            self.ros_mirobotstate = sensor_data.get('mirobotstate')
            self.ros_astra = sensor_data.get('astra')
            self.ros_emergency = sensor_data.get('emergency')
            self.ros_controlMode = sensor_data.get('controlMode')

        self.ros_imuData = recv_str.get('imuData')
        self.ros_odomData = recv_str.get('odomData')
        self.ros_batteryData = recv_str.get('batteryData')

    # #### 解析ros订阅的信息：电池电量，小车速度，角度，mirobot连接， 深度相机， 急停，控制模式等
    # def parse_data(self, recv_data):
    #     try:
            
    #         index = recv_data.find('}')
    #         if index != -1:
    #             first_part = recv_data[:index+1]
    #             second_part = recv_data[index+1:]

    #         if isinstance(second_part, str):  
    #             recv_str = json.loads(second_part)
            
    #         if 'sensorData_' in recv_str: 
    #             commandType = recv_str["sensorData_"]
    #             self.ros_battery = commandType.get('battery')
    #             self.ros_speed = commandType.get('speed')
    #             self.ros_angle = commandType.get('angle')
    #             self.ros_mirobotstate= commandType.get('mirobotstate')
    #             self.ros_astra = commandType.get('astra')
    #             self.ros_emergency = commandType.get('emergency')
    #             self.ros_controlMode = commandType.get('controlMode')

    #             self.ros_imuData = recv_str['imuData']
    #             self.ros_odomData = recv_str['odomData']
    #             self.ros_batteryData = recv_str['batteryData']

    #     except json.JSONDecodeError as e:  
    #         print(f"JSON解析错误: {e}")  


# if __name__ == '__main__':
#     rc = ROS_Client(server_ip='192.168.1.252', server_port=8888)
#     # rc = ROS_Client(server_ip='192.168.1.10', server_port=8527)
#     rc.connect()
#     # Example usage of sending commands to the server
#     # rc.send_command("ZBot", roscore="start", ros_run="zbot zbot_node.py", ros_topic="/zbot/odom", ros_service="/zbot/set_mode", ros_gcode="zbot_gcode.py")
#     # rc.send_command(ros_run="rosrun", ros_service="service")
#     gcode = "gcode: '\$h'"
#     rc.send_command(ros_run="rosrun",ros_gcode="GCode", data=gcode)

#     # rc.send_command( ros_launch="roslaunch", ros_bag="zbot3_drive", data="zbot3_bringup.launch")
#     rc.close()